//
// Created by syshen on 2021/6/29.
//

#include <cmath>
#include "NvInferPlugin.h"
#include "craft.h"
#include "utils/clipper.hpp"
#include "utils/RrConfig.h"
#include "utils/utils.h"

int DraftDet::serialization(const std::string &model_file, const std::string &engine_file) {
    Logger gLogger(nvinfer1::ILogger::Severity::kERROR);//kWARNING, kINFO, kINTERNAL_ERROR

    int device_count = 0;
    cudaGetDeviceCount(&device_count);
    assert(device_count > 0);
    if ((device_id + 1) > device_count) {
        device_id = 0;
    }
    std::cout << "GPU_nums: " << device_count << std::endl;
    std::cout << "Dbnet GPU ID: " << device_id << std::endl;
    cudaSetDevice(device_id);

    initLibNvInferPlugins(&gLogger, "");

//    std::unique_ptr <nvinfer1::IRuntime, nvinfer1::Destroy<nvinfer1::IRuntime>> runtime{
//            nvinfer1::createInferRuntime(gLogger)};

    std::unique_ptr <nvinfer1::IBuilder, nvinfer1::Destroy<nvinfer1::IBuilder>> builder{
            nvinfer1::createInferBuilder(gLogger)};
    if (!builder) {
        return -1;
    }

    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    std::unique_ptr <nvinfer1::INetworkDefinition, nvinfer1::Destroy<nvinfer1::INetworkDefinition>> network{
            builder->createNetworkV2(explicitBatch)};
    if (!network) {
        return -1;
    }

    std::unique_ptr <nvonnxparser::IParser, nvinfer1::Destroy<nvonnxparser::IParser>> parser{
            nvonnxparser::createParser(*network, gLogger)};
    if (!parser) {
        return -1;
    }

    std::unique_ptr <nvinfer1::IBuilderConfig, nvinfer1::Destroy<nvinfer1::IBuilderConfig>> config{
            builder->createBuilderConfig()};
    if (!config) {
        return -1;
    }

    auto parsed = parser->parseFromFile(model_file.c_str(), static_cast<int>(nvinfer1::ILogger::Severity::kERROR));
    if (!parsed) {
        std::cout << "ERROR: could not parse input engine." << std::endl;
        return -1;
    }

    builder->setMaxWorkspaceSize(((1ULL << 30) * 3));

    {
        builder->setFp16Mode(true);
        builder->setMaxBatchSize(2);
        auto profile = builder->createOptimizationProfile();
        profile->setDimensions(network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kMIN,
                               nvinfer1::Dims4{1, 3, 256, 256});
        profile->setDimensions(network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kOPT,
                               nvinfer1::Dims4{1, 3, 1280, 1280});
        profile->setDimensions(network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kMAX,
                               nvinfer1::Dims4{1, 3, 2560, 2560});
        config->addOptimizationProfile(profile);
        config->addOptimizationProfile(profile);
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
    }
    config->setMaxWorkspaceSize(static_cast<size_t>(16) << 20);
    //config->setFlag(nvinfer1::BuilderFlag::kDISABLE_TIMING_CACHE);
    std::unique_ptr <nvinfer1::ICudaEngine, nvinfer1::Destroy<nvinfer1::ICudaEngine>> serialize_eigine{builder->buildEngineWithConfig(*network, *config)};

    // 序列化
    std::unique_ptr <nvinfer1::IHostMemory, nvinfer1::Destroy<nvinfer1::IHostMemory>> trtModelStream{
            serialize_eigine->serialize()};
    nvinfer1::writeBuffer(trtModelStream->data(), trtModelStream->size(), engine_file.c_str());

    return 0;
}

int DraftDet::init(const std::string &model_file) {
    Logger gLogger(nvinfer1::ILogger::Severity::kERROR);//kWARNING, kINFO, kINTERNAL_ERROR
    int device_count = 0;
    cudaGetDeviceCount(&device_count);
    assert(device_count > 0);
    if ((device_id + 1) > device_count) {
        device_id = 0;
    }
    std::cout << "GPU_nums: " << device_count << std::endl;
    std::cout << "CRAFT GPU ID: " << device_id << std::endl;
    cudaSetDevice(device_id);

    //反序列化
    size_t dot_pos = model_file.find(".");
    std::string engine_file = model_file.substr(0, dot_pos) + ".engine";
    std::string buffer = nvinfer1::readBuffer(engine_file.c_str());
    if (trt_serialize) {

        if (!buffer.size()) {
            std::string Trtexec = RrConfigMgr::GetInstance().ReadString("ocrConfig", "Trtexec", "");
            std::string name = Trtexec + "/bin/trtexec --onnx=";
            name += model_file + " --tacticSources=-cublasLt,+cublas --workspace=1024 ";
            name += "--minShapes=input_0:1x3x256x256 --optShapes=input_0:1x3x1280x1280 --maxShapes=input_0:1x3x2560x2560 --shapes=input_0:1x3x1280x1280 ";
            name += "--fp16 --saveEngine=" + engine_file;
            system(name.c_str());
        }
    } else {
        serialization(model_file, engine_file);
    }
    std::unique_ptr <nvinfer1::IRuntime, nvinfer1::Destroy<nvinfer1::IRuntime>> mRuntime{
            nvinfer1::createInferRuntime(gLogger)};
    buffer = nvinfer1::readBuffer(engine_file.c_str());
    mEngine.reset(mRuntime->deserializeCudaEngine(buffer.data(), buffer.size(), nullptr));
    context.reset(mEngine->createExecutionContext());
    assert(mEngine->getNbBindings() == MAX_BINDING_NUMS);
    cudaStreamCreate(&stream_t);

    return 0;
}

bool DraftDet::getImgSize(const cv::Mat &tmp_, cv::Mat &output, float &ratio_x, float &ratio_y) {
    //图片以最大边对齐，缩放到1280的尺度上，保证长宽都是32的倍数
    cv::Mat tmp;
    cv::cvtColor(tmp_, tmp, cv::COLOR_BGR2RGB);
    int img_w = tmp.cols;
    int img_h = tmp.rows;
    int img_depth = tmp.channels();
    float ratio = 1280.f / std::max(img_h, img_w);
    int target_h = img_h * ratio;
    int target_w = img_w * ratio;
    cv::resize(tmp, tmp, cv::Size(target_w, target_h));
    tmp.convertTo(tmp, CV_32FC3);
    int target_h32 = std::ceil((float) target_h / 32) * 32;
    int target_w32 = std::ceil((float) target_w / 32) * 32;
    copyMakeBorder(tmp, output, 0, target_h32 - target_h, 0, target_w32 - target_w, cv::BORDER_CONSTANT,
                   cv::Scalar(0, 0, 0));
    channels = img_depth;
    inputH = target_h32;
    inputW = target_w32;
    ratio_x = (float) tmp_.cols / (float) target_w * 2;
    ratio_y = (float) tmp_.rows / (float) target_h * 2;
    return true;
}

bool DraftDet::Input(cv::Mat &src, float *data) {
    const float mean[3]{0.485f, 0.456f, 0.406f}; // Also in RGB order
    const float variance[3] = {0.229f, 0.224f, 0.225f};
    cv::Mat tmp_imt(inputH, inputW, CV_32FC3, cv::Scalar(0, 0, 0));

//    mPPMs.resize(1);
    int volImg = channels * inputH * inputW;
//    for (int r = 0; r < src.rows; ++r) {
//        for (int c = 0; c < src.cols; ++c) {
//            cv::Vec3f &pt = src.at<cv::Vec3f>(r, c);
//            cv::Vec3f &pt_img = tmp_imt.at<cv::Vec3f>(r, c);
//            pt_img.val[0] = (pt.val[0] - mean[0] * 255.f) / (variance[0] * 255.f);
//            pt_img.val[1] = (pt.val[1] - mean[1] * 255.f) / (variance[1] * 255.f);
//            pt_img.val[2] = (pt.val[2] - mean[2] * 255.f) / (variance[2] * 255.f);
//        }
//    }
////    float *hostDataBuffer = static_cast<float *>(buffers.getHostBuffer(mParams.inputTensorNames[0]));
////    mPPMs[0] = tmp_imt;
    for (int i = 0, volImg = channels * inputH * inputW; i < 1; ++i) {
        float *img_data = reinterpret_cast<float *>(src.data);
        for (int c = 0; c < channels; ++c) {
            for (unsigned j = 0, volChl = inputH * inputW; j < volChl; ++j) {
                data[i * volImg + c * volChl + j] = (img_data[j * channels + c]/255.0f -mean[c]) /(variance[c]) ;
            }
        }
    }

//    const float mean[3]{0.485f, 0.456f, 0.406f}; // Also in RGB order
//    const float variance[3] = {0.229f, 0.224f, 0.225f};
//    int i = 0;
//    for (int row = 0; row < src.rows; ++row) {
//        uchar *uc_pixel = src.data + row * src.step;
//        for (int col = 0; col < src.cols; ++col) {
//            data[i] = (uc_pixel[2] / 255.0f - mean[2]) / variance[2];
//            data[i + src.rows * src.cols] = (uc_pixel[1] / 255.0f - mean[1]) / variance[1];
//            data[i + 2 * src.rows * src.cols] = (uc_pixel[0] / 255.0f - mean[0]) / variance[0];
//            uc_pixel += 3;
//            ++i;
//        }
//    }

    return true;
}

bool get_mini_boxess(cv::RotatedRect &rotated_rect, cv::Point2f rect[],
                     int min_size) {

    cv::Point2f temp_rect[4];
    rotated_rect.points(temp_rect);
    for (int i = 0; i < 4; i++) {
        for (int j = i + 1; j < 4; j++) {
            if (temp_rect[i].x > temp_rect[j].x) {
                cv::Point2f temp;
                temp = temp_rect[i];
                temp_rect[i] = temp_rect[j];
                temp_rect[j] = temp;
            }
        }
    }
    int index0 = 0;
    int index1 = 1;
    int index2 = 2;
    int index3 = 3;
    if (temp_rect[1].y > temp_rect[0].y) {
        index0 = 0;
        index3 = 1;
    } else {
        index0 = 1;
        index3 = 0;
    }
    if (temp_rect[3].y > temp_rect[2].y) {
        index1 = 2;
        index2 = 3;
    } else {
        index1 = 3;
        index2 = 2;
    }

    rect[0] = temp_rect[index0];  // Left top coordinate
    rect[1] = temp_rect[index1];  // Left bottom coordinate
    rect[2] = temp_rect[index2];  // Right bottom coordinate
    rect[3] = temp_rect[index3];  // Right top coordinate

    if (rotated_rect.size.width < min_size ||
        rotated_rect.size.height < min_size) {
        return false;
    } else {
        return true;
    }
}

std::vector <std::vector<cv::Point2f>>
DraftDet::Output(const cv::Mat &txt, const cv::Mat &link, const cv::Mat &input, float &ratio_x, float &ratio_y) {

    //把output中的结果分开，即txt、link
    cv::Mat cvmat_text = txt * 255, cvmat_link = link * 255;
    cvmat_text.setTo(255, cvmat_text > 255);
    cvmat_text.setTo(0, cvmat_text < 0);
    cvmat_link.setTo(0, cvmat_link < 0);
    cvmat_link.setTo(255, cvmat_link > 255);
    cv::threshold(cvmat_text, cvmat_text, 0.4 * 255, 255, cv::THRESH_BINARY);
    cv::threshold(cvmat_link, cvmat_link, 0.4 * 255, 255, cv::THRESH_BINARY);
    //把txt、link二值化之后的像素对应位置相加
    cv::Mat cvmat_text_link = cvmat_text + cvmat_link;
    cvmat_text_link.setTo(255, cvmat_text_link > 255);
    cv::Mat seg_int;
    cvmat_text_link.convertTo(seg_int, CV_8UC1);
    //初始化找轮廓函数的输出
    cv::Mat labels;
    cv::Mat stats;
    cv::Mat centroids;
    int nccomps = cv::connectedComponentsWithStats(
            seg_int, //二值图像
            labels,     //和原图一样大的标记图
            stats, //nccomps×5的矩阵 表示每个连通区域的外接矩形和面积（就是pixel的个数）
            centroids, //nccomps×2的矩阵 表示每个连通区域的质心
            4
    );
    std::vector <std::vector<cv::Point2f>> boxes;//初始化输出

    for (int l = 1; l < nccomps; l++) {
        int area = stats.at<int>(l, cv::CC_STAT_AREA);
        if (area < 10) continue;
        double max_th;
        //仅赋值第l个轮廓的值到seg，seg里只有一个轮廓（一个词条）
        cv::Mat seg = labels == l;
        cv::Mat seg_;
        seg.convertTo(seg_, CV_32FC1);
        seg_ = (txt * 255) & seg_;
        cv::minMaxLoc(seg_, nullptr, &max_th);
        seg.convertTo(seg, CV_8UC1);
        seg *= 255;
        if (max_th < 0.7f * 255)
            continue;
        //确定轮廓区域，对指定区域进行形态学操作
        int x = stats.at<int>(l, cv::CC_STAT_LEFT);
        int y = stats.at<int>(l, cv::CC_STAT_TOP);
        int w = stats.at<int>(l, cv::CC_STAT_WIDTH);
        int h = stats.at<int>(l, cv::CC_STAT_HEIGHT);
        int niter = int(sqrtf(area * std::min(w, h) / (w * h)) * 2);
        int sx = x - niter;
        int ex = x + w + niter + 1;
        int sy = y - niter;
        int ey = y + h + niter + 1;
        sx = sx < 0 ? 0 : sx;
        sy = sy < 0 ? 0 : sy;
        ex = ex > inputW / 2 - 1 ? inputW / 2 - 1 : ex;
        ey = ey > inputH / 2 - 1 ? inputH / 2 - 1 : ey;
        cv::Rect rect;
        rect.x = sx;
        rect.y = sy;
        rect.height = ey - sy;
        rect.width = ex - sx;
        cv::Mat structure_element = cv::getStructuringElement(cv::MORPH_RECT,
                                                              cv::Size(niter, niter)); //设置膨胀/腐蚀的核为矩形，大小为3*3
        cv::dilate(seg(rect), seg(rect), structure_element); //膨胀
        //查找轮廓
        std::vector <std::vector<cv::Point>> contours;
        std::vector <cv::Vec4i> hierarcy;
        findContours(seg, contours, hierarcy, cv::RETR_CCOMP, cv::CHAIN_APPROX_SIMPLE);
#ifdef DEBUG
        std::vector <cv::Rect> boundRect(contours.size());  //定义外接矩形集合
        std::vector <cv::RotatedRect> box(contours.size()); //定义最小外接矩形集合
        cv::Point2f rect_4[4];
#endif
        cv::Point2f rect_re[4];
        cv::Point2f order_rect[4];
        for (int i = 0; i < contours.size(); i++) {
            std::vector <cv::Point2f> points;
            cv::RotatedRect rotated_rect = cv::minAreaRect(cv::Mat(contours[i]));
            if (!get_mini_boxess(rotated_rect, rect_re, 5)) {
                std::cout << "box too small" << std::endl;
                continue;
            }
            for (int k = 0; k < 4; k++) {
                order_rect[k] = rect_re[k];
                order_rect[k].x = int(order_rect[k].x * ratio_x);
                order_rect[k].y = int(order_rect[k].y * ratio_y);
                points.push_back(order_rect[k]);
            }
            if (points.size() > 0) boxes.push_back(points);
#ifdef DEBUG
            box[i] = minAreaRect(cv::Mat(contours[i]));  //计算每个轮廓最小外接矩形
            boundRect[i] = boundingRect(cv::Mat(contours[i]));
            box[i].points(rect_4);  //把最小外接矩形四个端点复制给rect数组
            rectangle(input, cv::Point(boundRect[i].x * ratio_x, boundRect[i].y * ratio_y),
                      cv::Point((boundRect[i].x + boundRect[i].width) * ratio_x,
                                (boundRect[i].y + boundRect[i].height) * ratio_y),
                      cv::Scalar(0, 255, 0), 2, 8);
#endif
        }
    }
#ifdef DEBUG
    cv::imwrite("/home/haier/Templates/wqg/code/ocr/general-ocr/18.jpg", input);
#endif
    return boxes;
}

std::vector <std::vector<cv::Point2f>>
DraftDet::Output(const cv::Mat &txt, const cv::Mat &link, float &ratio_x, float &ratio_y) {
    //把output中的结果分开，即txt、link
    cv::Mat cvmat_text = txt * 255, cvmat_link = link * 255;
    cvmat_text.setTo(255, cvmat_text > 255);
    cvmat_text.setTo(0, cvmat_text < 0);
    cvmat_link.setTo(0, cvmat_link < 0);
    cvmat_link.setTo(255, cvmat_link > 255);
    cv::threshold(cvmat_text, cvmat_text, 0.4 * 255, 255, cv::THRESH_BINARY);
    cv::threshold(cvmat_link, cvmat_link, 0.4 * 255, 255, cv::THRESH_BINARY);
    //把txt、link二值化之后的像素对应位置相加
    cv::Mat cvmat_text_link = cvmat_text + cvmat_link;
    cvmat_text_link.setTo(255, cvmat_text_link > 255);
    cv::Mat seg_int;
    cvmat_text_link.convertTo(seg_int, CV_8UC1);
    //初始化找轮廓函数的输出
    cv::Mat labels;
    cv::Mat stats;
    cv::Mat centroids;
    int nccomps = cv::connectedComponentsWithStats(
            seg_int, //二值图像
            labels,     //和原图一样大的标记图
            stats, //nccomps×5的矩阵 表示每个连通区域的外接矩形和面积（就是pixel的个数）
            centroids, //nccomps×2的矩阵 表示每个连通区域的质心
            4
    );
    std::vector <std::vector<cv::Point2f>> boxes;//初始化输出

    for (int l = 1; l < nccomps; l++) {
        int area = stats.at<int>(l, cv::CC_STAT_AREA);
        if (area < 10) continue;
        double max_th;
        //仅赋值第l个轮廓的值到seg，seg里只有一个轮廓（一个词条）
        cv::Mat seg = labels == l;
        cv::Mat seg_;
        seg.convertTo(seg_, CV_32FC1);
        seg_ = (txt * 255) & seg_;
        cv::minMaxLoc(seg_, nullptr, &max_th);
        seg.convertTo(seg, CV_8UC1);
        seg *= 255;
        if (max_th < 0.7f * 255)
            continue;
        //确定轮廓区域，对指定区域进行形态学操作
        int x = stats.at<int>(l, cv::CC_STAT_LEFT);
        int y = stats.at<int>(l, cv::CC_STAT_TOP);
        int w = stats.at<int>(l, cv::CC_STAT_WIDTH);
        int h = stats.at<int>(l, cv::CC_STAT_HEIGHT);
        int niter = int(sqrtf(area * std::min(w, h) / (w * h)) * 2);
        int sx = x - niter;
        int ex = x + w + niter + 1;
        int sy = y - niter;
        int ey = y + h + niter + 1;
        sx = sx < 0 ? 0 : sx;
        sy = sy < 0 ? 0 : sy;
        ex = ex > inputW / 2 - 1 ? inputW / 2 - 1 : ex;
        ey = ey > inputH / 2 - 1 ? inputH / 2 - 1 : ey;
        cv::Rect rect;
        rect.x = sx;
        rect.y = sy;
        rect.height = ey - sy;
        rect.width = ex - sx;
        cv::Mat structure_element = cv::getStructuringElement(cv::MORPH_RECT,
                                                              cv::Size(niter, niter)); //设置膨胀/腐蚀的核为矩形，大小为3*3
        cv::dilate(seg(rect), seg(rect), structure_element); //膨胀
        //查找轮廓
        std::vector <std::vector<cv::Point>> contours;
        std::vector <cv::Vec4i> hierarcy;
        findContours(seg, contours, hierarcy, cv::RETR_CCOMP, cv::CHAIN_APPROX_SIMPLE);
        cv::Point2f rect_re[4];
        cv::Point2f order_rect[4];
        for (int i = 0; i < contours.size(); i++) {
            std::vector <cv::Point2f> points;
            cv::RotatedRect rotated_rect = cv::minAreaRect(cv::Mat(contours[i]));
            if (!get_mini_boxess(rotated_rect, rect_re, 5)) {
                //std::cout << "box too small" << std::endl;
                continue;
            }
            for (int k = 0; k < 4; k++) {
                order_rect[k] = rect_re[k];
                order_rect[k].x = int(order_rect[k].x * ratio_x);
                order_rect[k].y = int(order_rect[k].y * ratio_y);
                points.push_back(order_rect[k]);
            }
            if (points.size() > 0) boxes.push_back(points);
        }
    }
    return boxes;
}

std::vector <std::vector<cv::Point2f>> DraftDet::extract(const cv::Mat &img_input) {
    cudaSetDevice(device_id);
    cv::Mat img_input_resize;
    float ratio_x, ratio_y;
    getImgSize(img_input, img_input_resize, ratio_x, ratio_y);

    int inputBufferSize = channels * inputH * inputW;
    nvinfer1::Dims4 inputDims{1, 3, inputH, inputW};
    CHECK(cudaMalloc(&bindings[0], inputBufferSize * sizeof(float)));
    int outputBufferSize = 2 * (inputH / 2) * (inputW / 2);
    CHECK(cudaMalloc(&bindings[1], outputBufferSize * sizeof(float)));

    float *data = new float[channels * inputH * inputW];
    float *output_ = new float[outputBufferSize];
    Input(img_input_resize, data);
    CHECK(cudaMemcpyAsync((float *) bindings[0], \
                data, inputBufferSize * sizeof(float), \
                cudaMemcpyHostToDevice, stream_t));

    context->setOptimizationProfile(0);
    context->setBindingDimensions(0, inputDims);
    bool status = context->enqueueV2(bindings, stream_t, nullptr);
    if (!status) {
        std::cout << "Enqueue failed" << std::endl;
    }
    cudaStreamSynchronize(stream_t);
    CHECK(cudaMemcpyAsync(output_, \
                (float *) bindings[1], \
                outputBufferSize * sizeof(float), cudaMemcpyDeviceToHost, stream_t));
    for (void *ptr : bindings) {
        cudaFree(ptr);
    }
    cv::Mat text(inputH / 2, inputW / 2, CV_32FC1, output_), linktext(inputH / 2, inputW / 2, CV_32FC1,
                                                                      output_ + outputBufferSize / 2);
#ifdef DEBUG
    std::vector <std::vector<cv::Point2f>> boxes = Output(text, linktext, img_input, ratio_x, ratio_y);
#else
    std::vector <std::vector<cv::Point2f>> boxes = Output(text, linktext, ratio_x, ratio_y);
#endif
    delete[] data;
    if (output_) {
        delete[] output_;
        output_ = nullptr;
    }
    return boxes;
}

DraftDet::~DraftDet() {
    cudaStreamDestroy(stream_t);
}

REGISTER_TEXT_DETECTION_CLASS(DraftDet);
